"""
    @brief:simulation environment
    @Editor:CJH
    @Date:2025/3/11
"""
import os.path
import sys
import threading
import time

sys.path.append("utils")
import cv2
import numpy as np
from gym import Env
from dataloader import Dataloader

font = cv2.FONT_HERSHEY_COMPLEX_SMALL
import matplotlib.pyplot as plt

start = time.time()


class Point(object):
    def __init__(self, name, x_max, x_min, y_max, y_min, icon_path, icon_shape=(50, 50)):
        self.x = 0
        self.y = 0
        self.x_min = x_min
        self.x_max = x_max
        self.y_min = y_min
        self.y_max = y_max
        self.name = name
        self.current_grid_pos = (0, 0)  # x,y
        self.icon = cv2.imread(f"{icon_path}")
        self.icon = cv2.resize(self.icon, icon_shape)

    def set_position(self, x, y):
        self.x = self.clamp(x, self.x_min, self.x_max - self.icon.shape[1])
        self.y = self.clamp(y, self.y_min, self.y_max - self.icon.shape[0])

    def get_position(self):
        return self.x, self.y

    def move(self, del_x, del_y, sim_):
        next_x = del_x + self.x
        next_y = del_y + self.y
        # 如果移动之后的距离超出了，返回False表明这次的移动是失败的
        if self.x_min < next_x < self.x_max - self.icon.shape[1] and self.y_min < next_y < self.y_max - self.icon.shape[
            0] and not sim_:
            self.x += del_x
            self.y += del_y
            return True
        elif self.x_min < next_x < self.x_max - self.icon.shape[1] and self.y_min < next_y < self.y_max - \
                self.icon.shape[
                    0] and sim_:
            return True
        else:
            return False

    def clamp(self, n, minn, maxn):
        return max(min(maxn, n), minn)


# 只作为构建一个动态地图环境，可随时更新，用于训练DQN模型
class Map(Env):
    def __init__(self, canvas_size=(800, 800), split_x_nums=20, split_y_nums=20,
                 total_stride=1000):  # split_nums将地图分割为不同大小得栅格地图
        super(Map, self).__init__()
        # 定义地图相关信息
        self.total_stride = total_stride
        self.ep_return = 0  # 回归经验值
        self.stride_left = self.total_stride  # 训练步长
        self.title = "Rectangle board"
        self.canvas_shape = (*canvas_size, 3)  # shape=(y,x)
        self.grid_map = np.zeros((split_y_nums, split_x_nums))  # shape=(y,x)
        # Define mass_fraction_map
        self.mass_fraction_viewer = Dataloader(self.grid_map.shape)
        self.canvas = None
        # ground truth map
        self.mass_fraction_map = None
        # Gaussian noisy map
        self.Gaussian_mass_map = None

        # Define elements present inside the environment
        self.element = None  # feichangniubi，可以添加多个机器人agent
        # 构造地图边界，确保机器人不会越界
        self.y_min = 0
        self.x_min = 0
        self.y_max = self.canvas_shape[0]
        self.x_max = self.canvas_shape[1]

    # 可视化地图初始化
    def map_init(self, csv_path=""):
        if os.path.exists(csv_path):
            self.mass_fraction_viewer.merge(csv_path)
            # 载入真实值的气体分数地图
            self.mass_fraction_map = self.mass_fraction_viewer.init_map
            # 载入模拟真实情况的高斯噪声地图
            self.Gaussian_mass_map = self.mass_fraction_viewer.noisy_map

        canvas = np.ones(self.canvas_shape, dtype=np.uint8) * 255
        # draw row line
        for row in range(1, self.grid_map.shape[0] + 1):
            cv2.line(canvas, (0, int(self.canvas_shape[0] * row / self.grid_map.shape[0])),
                     (self.canvas_shape[1], int(self.canvas_shape[0] * row / self.grid_map.shape[0])), (255, 0, 0), 1)

        for col in range(1, self.grid_map.shape[1] + 1):
            cv2.line(canvas, (int(self.canvas_shape[1] * col / self.grid_map.shape[1]), 0),
                     (int(self.canvas_shape[1] * col / self.grid_map.shape[1]), self.canvas_shape[0]), (0, 0, 255), 1)

        alpha = 0.3  # 第一幅图像的权重
        beta = 0.7  # 第二幅图像的权重
        gamma = 0.5  # 亮度调整
        mass_fraction_map = cv2.resize(self.mass_fraction_map, self.canvas_shape[1::-1])
        depth_image = cv2.normalize(mass_fraction_map, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
        depth_image = cv2.applyColorMap(depth_image, cv2.COLORMAP_JET)
        canvas = cv2.addWeighted(canvas, alpha, depth_image, beta,
                                 gamma)

        text = 'Stride Left: {} | Rewards: {}'.format(self.stride_left, self.ep_return)

        # Put the info on canvas
        cv2.putText(canvas, text, (10, 20), font,
                    1, (0, 0, 0), 1, cv2.LINE_AA)

        self.canvas=canvas


    def create_bot(self, name, icon_path, start_point=(2, 2)):
        icon_shape = (
            int((self.canvas_shape[1] / self.grid_map.shape[1]) * 0.3),
            int((self.canvas_shape[0] / self.grid_map.shape[0]) * 0.3))

        origin_step = (int(self.canvas_shape[1] / (2 * self.grid_map.shape[1]) - icon_shape[1] / 2),
                       int(self.canvas_shape[0] / (2 * self.grid_map.shape[0]) - icon_shape[0] / 2))

        # 实例化一个机器人本体
        self.element = Point(name, x_max=self.x_max, x_min=self.x_min, y_max=self.y_max, y_min=self.y_min,
                             icon_path=icon_path,
                             icon_shape=icon_shape)
        # 创建机器人的时候，设定初始位置
        if start_point[0] < self.grid_map.shape[0] and start_point[1] < self.grid_map.shape[0]:
            self.element.set_position(
                int(start_point[0] * self.canvas_shape[1] / self.grid_map.shape[1] + origin_step[0]),
                int(start_point[1] * self.canvas_shape[0] / self.grid_map.shape[0] + origin_step[1]))
            self.grid_map[start_point[1], start_point[0]] = 1  # 在栅格地图上赋予机器人位置信息
        else:
            self.element.set_position(int(self.canvas_shape[1] / self.grid_map.shape[1] + origin_step[0]),
                                      int(self.canvas_shape[0] / self.grid_map.shape[0] + origin_step[1]))
        self.element.current_grid_pos = start_point  # 为机器人附上栅格地图的位置

    # 获取机器人在栅格地图中位置的函数
    def upgrade_grid_position(self):
        current_state = self.element.get_position()
        current_state = int(current_state[0] * self.grid_map.shape[1] / self.canvas_shape[1]), int(
            current_state[1] * self.grid_map.shape[0] / self.canvas_shape[0])
        # 记录已经走过的路径
        self.grid_map[current_state[1], current_state[0]] = 1  # 在栅格地图上赋予机器人位置信息
        self.element.current_grid_pos = current_state  # 为机器人附上栅格地图的位置

    # 多线程展现画布
    def show(self):
        while True:
            cv2.imshow(self.title, self.canvas)
            cv2.waitKey(1)


# 构建一个动态状态，可随时更新地图，point移动，用于训练DQN模型
class State(Map):
    def __init__(self, canvas_size=(800, 800), split_x_nums=20, split_y_nums=20, total_stride=100):
        Map.__init__(self, canvas_size, split_x_nums, split_y_nums, total_stride)
        # 一个格子位移的最小距离,应该为格子尺寸
        self.box_min_step = np.array([int(self.canvas_shape[1] / self.grid_map.shape[1]),
                                      int(self.canvas_shape[0] / self.grid_map.shape[0])])
        self.action_dict = {"upper": self.box_min_step * np.array([0, -1]),
                            "lower": self.box_min_step * np.array([0, 1]),
                            "left": self.box_min_step * np.array([-1, 0]),
                            "right": self.box_min_step * np.array([1, 0]),
                            "up_right": self.box_min_step * np.array([1, -1]),
                            "up_left": self.box_min_step * np.array([-1, -1]),
                            "low_right": self.box_min_step * np.array([1, 1]),
                            "low_left": self.box_min_step * np.array([-1, 1])}
        self.point_list = []

    def draw_elements_on_canvas(self):
        # Init the canvas
        self.map_init()
        text = 'Stride Left: {} | Rewards: {}'.format(self.stride_left, self.ep_return)
        # Put the info on canvas
        cv2.putText(self.canvas, text, (10, 20), font,
                    1, (0, 0, 0), 1, cv2.LINE_AA)
        # Draw the heliopter on canvas
        # 将每个点画出来
        for index, point in enumerate(self.point_list):
            if index < len(self.point_list) - 1:
                cv2.line(self.canvas, self.point_list[index], self.point_list[index + 1], (0, 255, 0), 5)
            cv2.circle(self.canvas, point, 5, (0, 255, 0), -1)
        elem_shape = self.element.icon.shape
        x, y = self.element.x, self.element.y
        self.canvas[y: y + elem_shape[0], x:x + elem_shape[1]] = self.element.icon

    def take_actions(self, action, sim_):  # 给予一个sim_仿真符号，为了测试运动是否可行的同时，不更改机器人动作
        if not isinstance(action, str):
            action = list(self.action_dict.keys())[action]
        assert self.element is not None, "Please add element to execute following task"
        assert action in self.action_dict.keys(), "Invalid Action"
        self.point_list.append(self.element.get_position())

        move_flag = self.element.move(self.action_dict[action][0],
                                      self.action_dict[action][1], sim_)  # 返回布尔值，确定这一步是否是正常的
        if not sim_:
            self.upgrade_grid_position()
            # 移动机器人，减少可移动步数，如果是假设移动，则不减少步数
            self.stride_left -= 1

        return move_flag

    def step(self, name, action):
        if not isinstance(action, str):
            action = list(self.action_dict.keys())[action]
        assert self.element is not None, "Please add element to execute following task"
        assert action in self.action_dict.keys(), "Invalid Action"
        # parameter setting
        k1, k2, k3, k4, k5, k6 = 2, 5, 1, 0.1, 0.2, 2  # 设置权重影响分数
        stride_ratio = self.stride_left / self.total_stride
        visited_ratio = self.grid_map.sum() / (self.grid_map.shape[0] * self.grid_map.shape[1])  # 时间因素

        max_mass_thresh = np.max(self.mass_fraction_map)
        done_s = False  # 定义这次step的是否完成的标志
        # 当前位置
        current_pos = self.element.current_grid_pos
        # 当前位置的气体质量分数值
        current_this_point_mass = self.mass_fraction_map[current_pos[1], current_pos[0]]
        # 移动前的位置和
        before_point_sum = self.grid_map.sum()

        if not self.take_actions(action, sim_=False):
            # pass
            self.ep_return -= 1  # 做靠近地图边界的选择，扣分

        # -------------------------------------after_taking_actions-----------------------------------
        # 移动后的位置
        next_pos = self.element.current_grid_pos
        # 移动后的位置的气体质量分数值
        next_this_point_mass = self.mass_fraction_map[next_pos[1], next_pos[0]]
        # 移动后的位置和
        after_point_sum = self.grid_map.sum()
        # 1.转换为ndarray数组输出的地图，表示机器人当前位置
        pos_grid_map = np.zeros(self.grid_map.shape)
        pos_grid_map[next_pos[1], next_pos[0]] = 1
        # 2.转换为ndarray数组输出的地图，表示质量分数地图
        quality_scores_map = np.array(self.grid_map * self.mass_fraction_map)

        # DQN policy
        def policy_(done_s):
            # 1.Exploration Reward
            visited_reward = k3 * (1 - visited_ratio)  # 初始 5，随访问增加减小
            exploration_bonus = stride_ratio * (
                visited_reward if after_point_sum - before_point_sum != 0 else 0)  # 如果是没探索过的区域，则给予奖励，随时间变化这个影响会越来越少

            # 2.Dynamic Probability Reward 来回震荡，并且离阈值差的很远则扣分
            dynamic_Probability = -1 if after_point_sum - before_point_sum == 0 and abs(
                next_this_point_mass - max_mass_thresh) > k4 * 1.5 else 0

            # 3.New height bonus
            # print(abs(next_this_point_mass-np.max(quality_scores_map)))
            known_max_pos = np.unravel_index(np.argmax(quality_scores_map), quality_scores_map.shape)  # y,x
            if next_pos == known_max_pos[::-1] and after_point_sum - before_point_sum != 0:  # 只要不重复返回刷分就行
                distance_reward = np.max(quality_scores_map) / max_mass_thresh  # 如果发现一个浓度更高的地方，则奖励，归一化奖励值
            else:
                distance_reward = 0

            # 4.Prior_map approach strategy
            # Prior_known_max_pos = np.unravel_index(np.argmax(self.Gaussian_mass_map), self.Gaussian_mass_map.shape)
            Gaussian_mass_map = self.Gaussian_mass_map / np.max(self.Gaussian_mass_map)
            Prior_distance_reward = Gaussian_mass_map[next_pos[1], next_pos[0]]
            # Prior_distance_reward = 1 / (np.linalg.norm(
            #     np.array(Prior_known_max_pos[::-1]) - next_pos) + 1e-6)  # 距离越小奖励越高
            reward = k5 * exploration_bonus + \
                     k6 * dynamic_Probability + \
                     k2 * distance_reward + \
                     k1 * Prior_distance_reward + \
                     - 0.1  # 步数惩罚

            # 最大值阈值奖励
            if abs(next_this_point_mass - max_mass_thresh) < k4:
                self.ep_return = 100
                done_s = True

                return reward, done_s

            return reward, done_s

        reward, done_s = policy_(done_s)
        self.ep_return += reward if not done_s else 0

        # taking_actions

        if self.stride_left <= 0 or done_s:
            print('\n', "本轮得分:", '\n', self.ep_return)
            plt.clf()
            plt.imshow(quality_scores_map / np.max(quality_scores_map), cmap='jet', vmin=0,
                       vmax=max_mass_thresh)  # cmap 指定颜色映射
            plt.colorbar(label='Value')  # 添加颜色条
            plt.title('Heatmap Example')  # 设置标题
            plt.xlabel('X-axis')  # X 轴标签
            plt.ylabel('Y-axis')  # Y 轴标签
            # plt.savefig(f"data/runing_record_map/{round(self.ep_return, 3)}.png")
            done_s = True  # 表明状态已经完成
        # 累计策略返回的分数

        return np.stack([pos_grid_map, quality_scores_map / np.max(quality_scores_map),
                         self.Gaussian_mass_map / np.max(self.Gaussian_mass_map)]), reward, done_s

    def render(self, mode="parallel"):
        assert mode in ["parallel", "step"], "Invalid mode, must be either \"parallel\" or \"step\""
        assert self.mass_fraction_map is not None
        if mode == "parallel":
            # 渲染的时候，直接开个新的线程进行并行渲染。
            t1 = threading.Thread(target=self.show, args=(), daemon=True)
            # 设置守护线程,即主线程启动才会启动子线程
            t1.start()
        elif mode == "step":
            pass

    def reset(self):
        self.ep_return = 0  # 回归经验值重置
        self.stride_left = self.total_stride  # 训练步长重置
        self.grid_map = np.zeros(self.grid_map.shape)  # 栅格概率地图重置
        self.point_list = []

    def close(self):
        sys.exit()


if __name__ == "__main__":
    env = State(canvas_size=(1500, 900), split_x_nums=30, split_y_nums=45)
    env.map_init(r"D:\Gas_detector\data\gas_mass_data\export9.csv")
    env.render(mode="parallel")
    counter = 0
    env.create_bot("sb1", r"D:\Gas_detector\1.jpg", (12, 23))

    while True:
        # input("go")
        env.draw_elements_on_canvas()
        cv2.imshow("mass", cv2.resize(env.mass_fraction_map, (600, 900)))
        key = cv2.waitKey(30)
        if key == 119:
            next_state, reward, done = env.step("sb1", 'upper')
            # print(reward, done)
        elif key == 115:
            next_state, reward, done = env.step("sb1", 'lower')
            # print(reward, done)
        elif key == 100:
            next_state, reward, done = env.step("sb1", 'right')
            # print(reward, done)
        elif key == 97:
            next_state, reward, done = env.step("sb1", 'left')
            # print(reward, done)
        elif key == 113:
            env.reset()
        # counter += 1
        # time.sleep(0.002)
    env.close()
